black wrapped the 122-char except clause, moving Exception onto its own
line while the disable comment stayed on the closing paren, so pylint's
W0718 no longer saw it. Shorten the line instead.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

for more information, see https://pre-commit.ci

… dim checks and cleanups

- check_gemm_dims: restore assert_dim_for_fp8_exec semantics (per-tensor
  leading%8 / last%16, out_features%8 not %16); rich error messages with
  dims on the eager path, constant torch._check messages under compile
  (Dynamo forbids tensor closures in _check message lambdas).
- test_te_linear_dynamic_shapes: the recompile assertion compared a
  nonexistent counter (always 0==0); use stats/unique_graphs and absorb
  the one-time lazy is_fsdp2 hasattr-guard recompile with a warmup.
- custom_op: None-sentinel dtype uint8 -> complex32; a genuinely empty
  FP8 uint8 buffer (batch=0) decoded as None and broke compilation.
- OpaqueValueBundle: type-tag _to_hashable (list/tuple/Size no longer
  compare equal), guard __getattr__ against copy/pickle recursion on
  underscored probes, render non-finite floats evaluably in __fx_repr__.
- Linear.forward: fetch the cuBLAS workspace only after the eager-fallback
  decision; explicit torch._dynamo.graph_break(msg=...) so fullgraph=True
  errors carry the fallback reason instead of breaking on warnings.warn.
- warn_compile_disabled: move the 'use a newer PyTorch build' advice to
  the version-related call sites only.
- Comment/docstring/typography/pylint-disable cleanups in custom_op;
  test cosmetics (use_compile arg name, argparse-time validation of
  --compile/--use-cuda-graphs, merged NVINSPECT skips, docstring fixes);
  export get_cublas_workspace from cpp_extensions.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

… eager dim asserts

- check_gemm_dims is now a compile-only torch._check guard emitter, called
  from the compiled-op branch; eager dim validation returns to the op impl
  (assert + assert_dim_for_fp8_exec, as on main) so eager pays no overhead
  and keeps full error messages with dims.
- Trim verbose test docstrings/comments (te.Linear section, warmup helper,
  cudagraph-skip helper); describe the dynamic-shape scope (leading dims)
  instead of the fix history.
- Drop the stale 'FP8 with symbolic shapes unsupported' comments: FP8 with a
  mark_dynamic batch works on current nightly (verified: one graph reused
  across batch sizes, numerics match eager).

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

…wo float8 reprs

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

_sp_out_leading/_sp_inp_leading -> _out_leading_from_inp/_inp_leading_from_out;
shorten the weight_workspace field comment; drop the to_tensor_spec caveat
paragraph.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

…navailable

PG_REFERENCE_OPAQUE is computed once at import (Dynamo-friendly constant);
compile_unsupported_reason reports a tp_group it cannot carry instead of the
misleading _UnsupportedAdapter TypeError at trace time.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

… module-docstring duplication

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

copyreg.pickle is process-wide: with the reducer installed, torch.save of
any object graph reaching a ProcessGroup silently succeeded and the
checkpoint failed only at torch.load (the reconstruct stub raises).
Restore the loud failure at save time; the cost is that inductor bypasses
the FX disk cache for compiled distributed graphs (with its own warning)
until the cache-key pickler handles real opaque objects upstream.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

DelayedScaling quantizers are not value-opaque, so the compiled path falls
back to eager, which errors under fullgraph=True.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

… UB test to FP8

Under ub_overlap_rs_dgrad the impl returns the plain high-precision
reduce-scatter output as dgrad (the grad_input_quantizer only feeds the
communication buffer), while the fake declared a quantized dgrad spec --
an op output-contract mismatch.

test_linear_with_overlap_compile now also runs fp8_current_scaling and
mxfp8 for the column-parallel cases (bulk and DGRAD+RS); FP8 row-parallel
stays skipped (forced differentiable fp8_output is unsupported under
compile) and delayed scaling is excluded like elsewhere.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

…e_final_attempt

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

# Conflicts:
#	tests/pytorch/test_torch_compile.py

Tests merged from main leave pending delayed-scaling amax reductions in
FP8GlobalStateManager; a later autocast __exit__ then calls raw tex
bindings, graph-breaking the fullgraph=True Linear tests.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

Parse the args dataclass's annotations once, at registration, into an
immutable _ArgPlan: per-field _FieldPlan records (_FieldKind + schema
slots) plus the derived layout -- schema string, slot order, gradient
placement, tensor-or-quantized offsets -- with duplicate-slot-name
validation. pack/unpack interpret the plan on each call. Replaces the
adapter classes and the four layout helpers that each re-walked them;
the op schema and Linear semantics are unchanged.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

Parse the fwd fake-impl result into a per-trace _OutputPlan (logical
outputs / saved tensors with their flat Tensor[] ranges) and use it as
the single structure behind forward_fn, setup_context and backward.
Backward now slices grads per user output from the plan stashed on ctx:
a grad_outputs field on the backward args receives the whole tuple,
otherwise grad_output receives the first output's grad -- removing the
flat_grads[0] single-output assumption. Also reject unions mixing
tensor types with other members at registration.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

tq offsets join grad targets as on-demand derivations from fields; the
shared-bundle slot presence is implied by the packed dict itself; only
the output ranges (not the whole output plan, which references specs
and their quantizers) are stashed on ctx for backward.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

…e_final_attempt

…ale dtype

- fused_mla_q_uproj: use _linear_backward_impl (renamed in this branch)
- swizzle_scales_for_gemm: allocate swizzled scale buffers as uint8 so the
  python-visible scale_inv dtype matches quantizer allocations (the compiled
  op's fake declares uint8; the e4m3-dtyped buffer broke NVFP4 under
  torch.compile on Blackwell)
- silence pylint false positive on type.__new__ via attribute

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

The cached weight workspace is updated in place on the first microbatch,
which the functional custom op (mutates_args=()) can't express. Covers
skip_fp8_weight_update too (it implies caching).

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

…e effects in fallback cases

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

The fwd fake declared (inp.shape[0], out_features) while the real impl views
the output to (1, out_features) for 1D inputs. The bwd impl rederives the
input shape from grad_output, which cannot recover rank-1, so the autograd
glue now stashes the true input shapes on ctx (SymInt-safe) and views the
returned grads back to them.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

Weight caching now falls back to eager under torch.compile, so with the
argument the xfail test never reached the train/eval switching it covers.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

A quantized dgrad can't cross the op boundary (grads are packed one plain
Tensor[] slot each), so AOT tracing crashed instead of falling back.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

An op with an unused output is legally DCE'd (mutates_args=()), silently
dropping its collectives and state updates. Register the ops in FX's
side-effect registry by default; NVTE_COMPILE_OP_SIDE_EFFECTS selects
token (ordered effect tokens, blocks reordering too, but incompatible
with cudagraph trees today) or 0 (off).

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

An allocated transpose may be stale (invalidated by _reset_caches);
reconstruction derived validity from presence and silently revalidated
it, so the compiled path could consume a pre-update transpose.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

- skip under fake tensors/modes (a fake workspace poisoned the
  process-global cache)
- also preallocate in _apply, covering .to()/.cuda()/to_empty() flows
  (meta-device init never ran the reset_parameters path)
- preallocate the UB workspace only when comm overlap is actually
  enabled, not on ub_name alone (MHA sets it unconditionally)
- drop the stream-capture assert: it broke previously-working eager
  CUDA-graph captures, while the cudagraph-trees hazard it aimed at
  (warmup-pool allocation) never triggers it anyway

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

A graph break inside try/finally cannot build a resume function, so
Dynamo marked the shared Linear.forward code object SKIP: one
unsupported config silently reverted every te.Linear in the process to
eager. Hoist the config checks into _compile_eager_fallback_reason and
exit through a dynamo-disabled eager re-entry; only quantizer-dependent
conditions remain in the late check.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>